{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Image to Image Inpainting Stable Diffusion\n",
    "\n",
    "Similar to the standard stable diffusion inpainting example, except with the addition of an `inner_image` argument.\n",
    "\n",
    "`image`, `inner_image`, and `mask` should have the same dimensions. `inner_image` should have an alpha (transparency) channel.\n",
    "\n",
    "The aim is to overlay two images, then mask out the boundary between `image` and `inner_image` to allow stable diffusion to make the connection more seamless. For example, this could be used to place a logo on a shirt and make it blend seamlessly.This script was contributed by [Alex McKinney](https://github.com/vvvm23) and the notebook by [Parag Ekbote](https://github.com/ParagEkbote)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pip install diffusers torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0b56a945eb5145598c4fd153bc658786",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pipeline components...:   0%|          | 0/7 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "An error occurred while trying to fetch /home/zeus/.cache/huggingface/hub/models--stable-diffusion-v1-5--stable-diffusion-inpainting/snapshots/8a4288a76071f7280aedbdb3253bdb9e9d5d84bb/unet: Error no file named diffusion_pytorch_model.safetensors found in directory /home/zeus/.cache/huggingface/hub/models--stable-diffusion-v1-5--stable-diffusion-inpainting/snapshots/8a4288a76071f7280aedbdb3253bdb9e9d5d84bb/unet.\n",
      "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "An error occurred while trying to fetch /home/zeus/.cache/huggingface/hub/models--stable-diffusion-v1-5--stable-diffusion-inpainting/snapshots/8a4288a76071f7280aedbdb3253bdb9e9d5d84bb/vae: Error no file named diffusion_pytorch_model.safetensors found in directory /home/zeus/.cache/huggingface/hub/models--stable-diffusion-v1-5--stable-diffusion-inpainting/snapshots/8a4288a76071f7280aedbdb3253bdb9e9d5d84bb/vae.\n",
      "Defaulting to unsafe serialization. Pass `allow_pickle=False` to raise an error instead.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cfa34d5822784b449f9014f3f1b0e4ef",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/50 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "import requests\n",
    "from PIL import Image\n",
    "from io import BytesIO\n",
    "from diffusers import DiffusionPipeline\n",
    "\n",
    "# Correct image URLs\n",
    "image_url = \"https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png\"\n",
    "inner_image_url = \"https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png\"\n",
    "mask_url = \"https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png\"\n",
    "\n",
    "# Function to load image from URL\n",
    "def load_image(url, mode=\"RGB\"):\n",
    "    response = requests.get(url)\n",
    "    if response.status_code == 200:\n",
    "        return Image.open(BytesIO(response.content)).convert(mode).resize((512, 512))\n",
    "    else:\n",
    "        raise FileNotFoundError(f\"Could not retrieve image from {url}\")\n",
    "\n",
    "# Load images\n",
    "init_image = load_image(image_url, mode=\"RGB\")\n",
    "inner_image = load_image(inner_image_url, mode=\"RGBA\")\n",
    "mask_image = load_image(mask_url, mode=\"RGB\")\n",
    "\n",
    "# Load the pipeline\n",
    "pipe = DiffusionPipeline.from_pretrained(\n",
    "    \"stable-diffusion-v1-5/stable-diffusion-inpainting\",\n",
    "    custom_pipeline=\"img2img_inpainting\",\n",
    "    torch_dtype=torch.float16\n",
    ")\n",
    "pipe = pipe.to(\"cuda\")\n",
    "\n",
    "# Inpainting\n",
    "prompt = \"a mecha robot sitting on a bench\"\n",
    "image = pipe(prompt=prompt, image=init_image, inner_image=inner_image, mask_image=mask_image).images[0]\n",
    "\n",
    "image.save(\"output.png\")\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
